from app_config import get_pro
import pandas as pd


def get_data(trade_date):
    df = pro.daily_basic(ts_code='',
                         trade_date=trade_date,
                         fields='ts_code,trade_date,close,turnover_rate,turnover_rate_f,total_share,float_share,'
                                'free_share,total_mv,'
                                'circ_mv')
    return df


if __name__ == '__main__':
    pro = get_pro()
    # 示例 DataFrame
    data = {
        'ts_code': ['600160.SH', '688728.SH', '601298.SH', '002463.SZ', '688072.SH', '002422.SZ', '600482.SH',
                    '002600.SZ', '688506.SH', '002032.SZ', '688047.SH', '000039.SZ', '603605.SH', '601061.SH']
    }
    start_date_str = '20231208'
    end_date_str = '20240531'

    ts_codes = pd.DataFrame(data)
    print(ts_codes)

    data = pro.trade_cal(exchange='SSE', is_open='1', start_date='20240101', end_date='20240601', fields='cal_date')
    data.sort_values(by=['cal_date'], inplace=True)
    first_date = data['cal_date'].iloc[0]
    last_date = data['cal_date'].iloc[-1]

    daily_first = get_data(first_date)
    print(first_date)
    ts_codes = pd.merge(ts_codes, daily_first[["ts_code", "close"]], on='ts_code', how='left')
    ts_codes.rename(columns={'close': first_date}, inplace=True)

    daily_end = get_data(last_date)
    print(last_date)
    ts_codes = pd.merge(ts_codes, daily_end[["ts_code", "close"]], on='ts_code', how='left')
    ts_codes.rename(columns={'close': last_date}, inplace=True)

    stock_basic = pro.stock_basic(exchange='', list_status='L', fields='ts_code,name')
    ts_codes = pd.merge(ts_codes, stock_basic, on='ts_code', how='left')

    ts_codes["up_down"] = (ts_codes[last_date] - ts_codes[first_date]) / ts_codes[first_date] * 100

    ts_codes.to_excel("price.xlsx", index=False)

    print(ts_codes)
